"""
利用retdec进行数据处理，获得functionSim的模型输入
"""
import os
from tqdm import tqdm
import shelve
import datetime
import subprocess
import signal
import json


def set_timeout(num, callback):
    """
        函数超时设置，避免一直卡在一个函数中
    """
    def wrap(func):
        # 收到信号 SIGALRM 后的回调函数，第一个参数是信号的数字，第二个参数是the interrupted stack frame.
        def handle(signum, frame):
            raise RuntimeError

        def to_do(*args, **kwargs):
            try:
                signal.signal(signal.SIGALRM, handle)  # 设置信号和回调函数
                signal.alarm(num)  # 设置 num 秒的闹钟
                print('start alarm signal.')
                r = func(*args, **kwargs)
                print('close alarm signal.')
                signal.alarm(0)  # 关闭闹钟
                return r
            except RuntimeError as e:
                callback()
        return to_do
    return wrap


def after_timeout():  # 超时后的处理函数
    print("Time out!")


timeOut = 600

logPath = r"/home/cyw/projects/function_sim_project/disassemblyTools/retdec/retdec_logs.txt"


def logs(inf):
    with open(logPath, "a+") as file:
        file.write(str(datetime.datetime.now())+"\t" + inf + "\n")


def getUnfinishedName(originPath, savePath):
    sampleList = {}
    finishList = {}
    unfinList = {}
    failedList = {}
    sLoc = 0
    fLoc = 0
    uLoc = 0
    # 获取所有的样本名
    for file in os.listdir(originPath):
        data = file.split(".")  # 找到样本文件
        if (len(data) == 1):
            sLoc += 1
            sampleList[data[0]] = True
    # 获取已经处理过的样本名
    for file in os.listdir(savePath):
        data = file.split(".")
        if (data[1] == "bak"):
            fLoc += 1
            finishList[data[0]] = True
    # 获得不能处理的样本
    path = "/home/cyw/projects/function_sim_project/disassemblyTools/retdec/failedSample.txt"
    with open(path) as file:
        for str in file.readlines():
            failedList[str.strip()] = True
    # 获得未被处理过的样本
    for name in sampleList.keys():
        if name not in finishList and name not in failedList:
            uLoc += 1
            unfinList[name] = True
    logs("样本总计{}个，已处理{}个,还剩{}个\n".format(sLoc, fLoc, uLoc))
    return unfinList


class retdecDevTools:
    dataTrans = ["mov", "push", "pop", "xchg", "in", "out", "xlat",
                 "cwd", "lea", "lds", "les", "lahf", "sahf", "pushf", "popf"]
    arithmetric = ["add", "adc", "inc", "daa", "sub", "sbb", "dec", "cmp",
                   "das", "mul", "imul", "div", "idiv", "xor", "or", "not", "sll", "srl"]
    controlTrans = ["call", "ret", "jmp", "loop", "loope", "loopne", "int", "into", "iret", "jz",
                    "je", "jnz", "jne", "jo", "jno", "jc", "jnc", "js", "jns", "jb", "jnae", "jnb", "jae", "ja",
                    "jnbe", "jna", "jbe", "jl", "jnge", "jnl", "jge", "jg", "jnle", "jng", "jle"]
    retdecPath = "/home/cyw/tools/retdec-install/bin/retdec-decompiler"
    success = False

    def __init__(self) -> None:
        pass

    def get_function_type_and_boundry(self, filePath):
        """
            判断函数类型和函数边界信息
            返回信息：
                {
                    functionName:type
                }
            local,dynamic import,static link
        """
        fncInf = {}
        trans = {}
        trans["staticallyLinked"] = "static link"
        trans["decompilerDefined"] = "local"
        trans["dynamicallyLinked"] = "dynamic import"
        file = open(filePath, 'r', encoding='utf-8')
        content = file.read()
        file.close()
        rs = json.loads(content)['functions']
        for func in rs:
            fncInf[func["name"]] = trans[func["fncType"]]
        return fncInf

    def gene_funtion_embedding(self, filePath):
        """
            利用反汇编代码，生成函数嵌入信息        
            提取函数的特征
                0：字符常量数量
                1：数值常量数量
                2：转移指令数量
                3：控制指令数量
                4：指令数量
                5：算术指令数量
                6：后代数量
                7：上一辈数量
            返回类型：
            {
                函数名：特征
            }

        """
        file = open(filePath, "r", encoding="utf-8")
        data = file.readlines()
        fncName = "---start-name----"
        resdict = {}
        temp = [0, 0, 0, 0, 0, 0, 0, 0]
        # 0：字符常量数量
        # 1：数值常量数量
        # 2：转移指令数量
        # 3：控制指令数量
        # 4：指令数量
        # 5：算术指令数量
        # 6：后代数量
        # 7：上一辈数量
        for line in data:
            if len(line) <= 12:
                continue
            elif line[:12] == "; function: ":
                if fncName != "---start-name----":
                    resdict[fncName] = temp.copy()
                fncName = line[12:].split(" ")[0]
                temp = [0, 0, 0, 0, 0, 0, 0, 0]
            elif line[:2] == "0x":  # 指令统计特征的提取
                temp[4] += 1
                instr = line[45:]
                instrName = instr.split(" ")[0]
                if instrName in self.dataTrans:
                    temp[2] += 1
                elif instrName in self.arithmetric:
                    temp[5] += 1
                elif instrName in self.controlTrans:
                    temp[3] += 1
                for i in instr.split(" "):
                    if i.isdecimal():
                        temp[1] += 1
                if '"' in instr:
                    temp[0] += 1
        resdict[fncName] = temp
        return resdict

    def obtain_call_graph_struct(self, filePath):
        """
            解析retdec生成的中间文件.c.cg.dot，
            提取生成图结构的信息
            返回值为：样本与下标的节点，x,y表示调用关系，同下标表示x调用y
        """
        file = open(filePath, "r", encoding="utf-8")
        data = file.readlines()
        ind = 0
        res = {}
        nameList = {}
        aList = []
        bList = []
        for line in data:
            if line.find('label="{') != -1:
                name = line[line.find('label="{')+8:].split("}")[0]
                nameList[name] = ind
                ind += 1
            elif line.find(" -> ") != -1:
                a = line[:line.find(" -> ")].strip()
                a = a[a.find("_")+1:]
                aList.append(a)
                b = line[line.find(" -> ")+4:]
                b = b[:-2]
                b = b[b.find("_")+1:]
                bList.append(b)
        res["name"] = nameList
        res["x"] = aList
        res["y"] = bList
        return res

    def gene_model_data(self, originPath, sampleName):
        """
            # cg
            # cgattr
            # functype
            # funcs_id
        """
        samplePath = originPath + "/" + sampleName
        res = {}
        funcs_id, cnt = {}, 0
        functype = {}
        try:
            cgInf = self.obtain_call_graph_struct(samplePath+".c.cg.dot")
            funcTypeOrigin = self.get_function_type_and_boundry(
                samplePath+".config.json")
            funcEmbedings = self.gene_funtion_embedding(samplePath+".dsm")
            print("正在处理数据~")
            # 为啥节点会对应不上===》这里取重复的节点好了,其中动态的无需函数嵌入
            # 初始化项目特征
            for name in funcTypeOrigin:
                if funcTypeOrigin[name] != "dynamic import":
                    if name in funcEmbedings and name in cgInf["name"]:
                        funcs_id[name] = cnt
                        functype[cnt] = funcTypeOrigin[name]
                        cnt += 1
                else:
                    if name in cgInf["name"]:
                        funcs_id[name] = cnt
                        functype[cnt] = funcTypeOrigin[name]
                        cnt += 1
            cgattr, cg = [], []
            for name in funcs_id.keys():
                cg.append([0]*cnt)
                if (functype[funcs_id[name]] != "dynamic import"):
                    cgattr.append(funcEmbedings[name])
                else:
                    cgattr.append([0, 0, 0, 0, 0, 0, 0, 0])
            removeEdgeNum = 0
            for i in range(len(cgInf["x"])):
                xName = cgInf["x"][i]
                yName = cgInf["y"][i]
                if xName in funcs_id and yName in funcs_id:
                    x = funcs_id[cgInf["x"][i]]
                    y = funcs_id[cgInf["y"][i]]
                    if cg[x][y] != 0:
                        print("重复 重复 重复")
                    cg[x][y] += 1
                else:
                    removeEdgeNum += 1
            nodeNum = len(cg)
            for num in range(nodeNum):
                for i in range(nodeNum):
                    if (cg[num][i] != 0):
                        cgattr[num][6] += 1
                for j in range(nodeNum):
                    if cg[j][num] != 0:
                        cgattr[num][7] += 1
            print("基本信息，节点数：{}，边数：{}，移除了{}条边".format(
                len(cg), len(cgInf["x"])-removeEdgeNum, removeEdgeNum))
            res["cg"] = cg
            res["cgattr"] = cgattr
            res["funcs_id"] = funcs_id
            res["functype"] = functype
            return res
        except Exception as e:
            path = "/home/cyw/projects/function_sim_project/disassemblyTools/retdec/failedSample.txt"
            with open(path, "a+") as file:
                file.write(sampleName+"\n")
            logs("retdec处理失败")
            logs("exception-->"+str(e))
            return False

    def save_data(self, savePath, res):
        # 保存数据
        file = shelve.open(savePath)
        file["cg"] = res["cg"]
        file["cgattr"] = res["cgattr"]
        file["functype"] = res["functype"]
        file["funcs_id"] = res["funcs_id"]
        file.close()
        print("样本保存成功")

    @set_timeout(timeOut, after_timeout)
    def analysisBinary(self, originPath, name):
        """
            进行样本的分析
            保存如下数据：
            # cg
            # cgattr
            # functypes
            # funcs_id
        """
        sample = originPath + "/" + name
        res = {}
        cmd = "/home/cyw/tools/retdec-install/bin/retdec-decompiler" + \
            " -s --backend-emit-cg "+sample
        self.success = False
        self.res = {}
        try:
            # 通过subprocess.run()函数运行Linux命令并获取输出结果
            result = subprocess.run(
                cmd, shell=True, capture_output=True, text=True)
            if result.returncode == 0:
                print("命令成功执行")
                print(result.stdout)
                self.success = True
            else:
                print("命令执行失败")
        except Exception as e:
            path = "/home/cyw/projects/function_sim_project/disassemblyTools/retdec/failedSample.txt"
            with open(path, "a+") as file:
                file.write(name+"\n")
            logs("retdec处理失败")
            logs("发生错误：{}".format(str(e)))

    def clean_temp_files(self, samplePath):
        """
            清除生成的临时文件
        """
        suffixs = [".bc", ".c", ".c.cg.dot", ".config.json", ".dsm", ".ll"]
        for name in suffixs:
            filename = samplePath+name
            if os.path.exists(filename):
                os.remove(filename)
                print(f"文件 {filename} 删除成功！")
            else:
                print(f"文件 {filename} 不存在。")

    def clean_Path_temp_files(self, originPath):
        sampleList = {}
        for file in os.listdir(originPath):
            data = file.split(".")  # 找到样本文件
            if (len(data) == 1):
                sampleList[data[0]] = True
        for name in tqdm(sampleList):
            self.clean_temp_files(originPath + "/"+name)
        print("中间文件已删除")

    # 分析文件夹下的样本，
    def analysisPath(self, originPath, savePath):
        unfinishedList = getUnfinishedName(originPath, savePath)
        logs("使用retdec,进行样本反汇编")
        for name in tqdm(unfinishedList):
            samplePath = originPath + "/" + name
            self.analysisBinary(originPath, name)
            if (self.success):
                res = self.gene_model_data(originPath, name)
                if res != False:
                    self.save_data(savePath+"/"+name, res)
            self.clean_temp_files(samplePath)
        logs("样本反汇编完成")


if __name__ == '__main__':
    originPath = r"/home/cyw/projects/data/functionsim/allorigindata"
    # originPath = r"/home/cyw/projects/data/functionsim/test"
    savePath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/retdec"
    retdecTool = retdecDevTools()
    # retdecTool.analysisPath(originPath,savePath)
    retdecTool.clean_Path_temp_files(originPath)
